/*
    Parallel Syscalls
    Rust version: @5mukx
*/


use ntapi::{ntldr::LDR_DATA_TABLE_ENTRY, ntpebteb::PEB};
use windows::{
    Wdk::Foundation::OBJECT_ATTRIBUTES,
    Win32::{
        Foundation::{CloseHandle, HANDLE, NTSTATUS, OBJ_CASE_INSENSITIVE, UNICODE_STRING},
        System::{
            Diagnostics::Debug::{IMAGE_NT_HEADERS64, IMAGE_SECTION_HEADER},
            IO::IO_STATUS_BLOCK,
            LibraryLoader::{GetModuleHandleW, GetProcAddress},
            Memory::{
                MEM_COMMIT, MEM_RESERVE, PAGE_EXECUTE_READWRITE, PAGE_READONLY, VirtualAlloc,
            },
            SystemServices::{IMAGE_DOS_HEADER, IMAGE_EXPORT_DIRECTORY},
            Threading::{GetCurrentProcess, INFINITE, Sleep},
        },
    },
    core::{Error, PCWSTR, PSTR, Result},
};

use ntapi::winapi::shared::ntdef::NT_SUCCESS;
use std::arch::asm;
use std::mem::{size_of, transmute};
use std::ptr::addr_of;
use std::ptr::{null, null_mut};
use std::slice::from_raw_parts;
use widestring::U16CString;

const MAX_SYSCALL_STUB_SIZE: usize = 64;
const MAX_NUMBER_OF_SYCALLS: usize = 1024;

// function pointers
type NtOpenFile = unsafe extern "system" fn(
    FileHandle: *mut HANDLE,
    DesiredAccess: u32,
    ObjectAttributes: *const OBJECT_ATTRIBUTES,
    IoStatusBlock: *mut IO_STATUS_BLOCK,
    ShareAccess: u32,
    OpenOptions: u32,
) -> NTSTATUS;

type NtCreateSection = unsafe extern "system" fn(
    SectionHandle: *mut HANDLE,
    DesiredAccess: u32,
    ObjectAttributes: *const OBJECT_ATTRIBUTES,
    MaximumSize: *const i64,
    SectionPageProtection: u32,
    AllocationAttributes: u32,
    FileHandle: HANDLE,
) -> NTSTATUS;

type NtMapViewOfSection = unsafe extern "system" fn(
    SectionHandle: HANDLE,
    ProcessHandle: HANDLE,
    BaseAddress: *mut *mut u8,
    ZeroBits: usize,
    CommitSize: usize,
    SectionOffset: *const i64,
    ViewSize: *mut usize,
    InheritDisposition: u32,
    AllocationType: u32,
    Win32Protect: u32,
) -> NTSTATUS;

type NtCreateThreadEx = unsafe extern "system" fn(
    ThreadHandle: *mut HANDLE,
    DesiredAccess: u32,
    ObjectAttributes: *const OBJECT_ATTRIBUTES,
    ProcessHandle: HANDLE,
    StartAddress: *mut u8,
    Parameter: *mut u8,
    Flags: u32,
    StackZeroBits: usize,
    SizeOfStackCommit: usize,
    SizeOfStackReserve: usize,
    BytesBuffer: *mut u8,
) -> NTSTATUS;

type RtlInitUnicodeString =
    unsafe extern "system" fn(DestinationString: *mut UNICODE_STRING, SourceString: PCWSTR);

// global syscall pointers
static mut NT_OPEN_FILE_CLEAN: Option<NtOpenFile> = None;
static mut NT_CREATE_SECTION_CLEAN: Option<NtCreateSection> = None;
static mut NT_MAP_VIEW_OF_SECTION_CLEAN: Option<NtMapViewOfSection> = None;

// syscall storage
static mut SYSCALL_NAMES: [String; MAX_NUMBER_OF_SYCALLS] =
    [const { String::new() }; MAX_NUMBER_OF_SYCALLS];
static mut SYSCALL_STUBS: [u64; MAX_NUMBER_OF_SYCALLS] = [0; MAX_NUMBER_OF_SYCALLS];

#[inline]
#[cfg(target_pointer_width = "64")]
fn get_peb() -> *const PEB {
    unsafe {
        let mut peb: *const PEB;
        asm!(
            "mov {0}, gs:[0x60]",
            out(reg) peb,
            options(nostack, nomem, preserves_flags)
        );
        peb
    }
}

fn build_syscall_stub(stub_region: u64, syscall_no: u32) -> u64 {
    let syscall_stub: [u8; 11] = [
        0x4c, 0x8b, 0xd1, // mov r10, rcx
        0xb8, 0x00, 0x00, 0x00, 0x00, // mov eax, xxx
        0x0f, 0x05, // syscall
        0xc3, // ret
    ];

    unsafe {
        std::ptr::copy_nonoverlapping(
            syscall_stub.as_ptr(),
            stub_region as *mut u8,
            syscall_stub.len(),
        );

        *(stub_region as *mut u8).add(4) = syscall_no as u8;
        *(stub_region as *mut u8).add(5) = (syscall_no >> 8) as u8;
        *(stub_region as *mut u8).add(6) = (syscall_no >> 16) as u8;
        *(stub_region as *mut u8).add(7) = (syscall_no >> 24) as u8;
    }

    stub_region
}

fn init_syscalls_from_ldrp_thunk_signature() -> bool {
    unsafe {
        let peb = get_peb();
        if peb.is_null() {
            #[cfg(feature = "debug")]
            println!("[!] PEB is null");
            return false;
        }

        // calculate the offset of InLoadOrderLinks in LDR_DATA_TABLE_ENTRY
        let ldr_entry_offset = {
            let entry: *const LDR_DATA_TABLE_ENTRY = std::ptr::null();
            let entry_addr = entry as usize;
            let links_addr = addr_of!((*entry).InLoadOrderLinks) as usize;
            links_addr - entry_addr
        };

        let ldr = (*peb).Ldr;
        let list_head = &(*ldr).InLoadOrderModuleList as *const _
            as *const ntapi::winapi::shared::ntdef::LIST_ENTRY;
        let mut cursor = (*list_head).Flink;

        let mut ntdll_base = 0u64;

        while cursor != list_head as _ {
            let entry = (cursor as usize - ldr_entry_offset) as *const LDR_DATA_TABLE_ENTRY;
            if (*entry).DllBase.is_null() {
                break;
            }
            let name = from_raw_parts(
                (*entry).BaseDllName.Buffer,
                (*entry).BaseDllName.Length as usize / 2,
            );

            let char_name: String = name.iter().map(|&c| c as u8 as char).collect();
            if char_name.to_lowercase() == "ntdll.dll" {
                ntdll_base = (*entry).DllBase as u64;
                println!("[+] Found ntdll.dll: Base=0x{:016x}", ntdll_base);
                break;
            }
            cursor = (*cursor).Flink;
        }

        if ntdll_base == 0 {
            #[cfg(feature = "debug")]
            println!("[!] Failed to find ntdll.dll");
            return false;
        }

        let dos_header = ntdll_base as *const IMAGE_DOS_HEADER;
        let nt_headers = (ntdll_base + (*dos_header).e_lfanew as u64) as *const IMAGE_NT_HEADERS64;

        let section_header = &(&(*nt_headers)).OptionalHeader as *const _ as *const u8;
        let section_header = section_header
            .add((*nt_headers).FileHeader.SizeOfOptionalHeader as usize)
            as *const IMAGE_SECTION_HEADER;

        println!(
            "[+] Found {} sections in ntdll.dll",
            (*nt_headers).FileHeader.NumberOfSections
        );

        let mut data_section_address = 0u64;
        let mut data_section_size = 0u32;

        for i in 0..(*nt_headers).FileHeader.NumberOfSections {
            let section = section_header.add(i as usize);
            let name_bytes = &(*section).Name[..8];
            let section_name = std::ffi::CStr::from_bytes_until_nul(name_bytes)
                .map(|cstr| cstr.to_str().unwrap_or("<invalid>"))
                .unwrap_or("<invalid>");
            println!(
                "[+] Section {}: Name={}, VirtualAddress=0x{:x}, VirtualSize=0x{:x}",
                i,
                section_name,
                (*section).VirtualAddress,
                (*section).Misc.VirtualSize
            );
            if section_name == ".data" {
                data_section_address = ntdll_base + (*section).VirtualAddress as u64;
                data_section_size = (*section).Misc.VirtualSize;
                println!(
                    "[+] Found .data section: Address=0x{:016x}, Size=0x{:x}",
                    data_section_address, data_section_size
                );
            }
        }

        if data_section_address == 0 || data_section_size < 16 * 5 {
            #[cfg(feature = "debug")]
            println!("[!] .data section not found or too small");
            return false;
        }

        let mut nt_open_file_num = 0u32;
        let mut nt_create_section_num = 0u32;
        let mut nt_map_view_of_section_num = 0u32;
        let signature = 0xb8d18b4c_u32;

        let mut found_offset = 0;
        for offset in 0..(data_section_size - 16 * 5) {
            let addr = (data_section_address + offset as u64) as *const u32;
            if *(addr) == signature
                && *(addr.add(4)) == signature
                && *(addr.add(8)) == signature
                && *(addr.add(12)) == signature
                && *(addr.add(16)) == signature
            {
                nt_open_file_num = *(addr.add(1));
                nt_create_section_num = *(addr.add(5));
                nt_map_view_of_section_num = *(addr.add(17));
                found_offset = offset;
                println!(
                    "[+] Found LdrpThunkSignature at offset 0x{:x}: NtOpenFile=0x{:x}, NtCreateSection=0x{:x}, NtMapViewOfSection=0x{:x}",
                    offset, nt_open_file_num, nt_create_section_num, nt_map_view_of_section_num
                );
                break;
            }
        }

        if nt_open_file_num == 0 || nt_create_section_num == 0 || nt_map_view_of_section_num == 0 {
            #[cfg(feature = "debug")]
            println!("[!] Failed to find syscall numbers");
            return false;
        }

        // dump first 128 bytes at the found offset
        let dump_addr = data_section_address + found_offset as u64;
        let bytes = from_raw_parts(dump_addr as *const u8, 128);
        print!("[+] First 128 bytes at 0x{:016x}:\n", dump_addr);
        for i in 0..8 {
            for j in 0..16 {
                print!("{:02x} ", bytes[i * 16 + j]);
            }
            print!("\n");
        }
        println!();

        let mem_pointer = VirtualAlloc(
            None,
            3 * MAX_SYSCALL_STUB_SIZE,
            MEM_RESERVE | MEM_COMMIT,
            PAGE_EXECUTE_READWRITE,
        ) as u64;

        if mem_pointer == 0 {
            #[cfg(feature = "debug")]
            println!("[!] VirtualAlloc for syscall stubs failed");
            return false;
        }

        NT_OPEN_FILE_CLEAN = Some(transmute(build_syscall_stub(mem_pointer, nt_open_file_num)));
        NT_CREATE_SECTION_CLEAN = Some(transmute(build_syscall_stub(
            mem_pointer + MAX_SYSCALL_STUB_SIZE as u64,
            nt_create_section_num,
        )));
        NT_MAP_VIEW_OF_SECTION_CLEAN = Some(transmute(build_syscall_stub(
            mem_pointer + 2 * MAX_SYSCALL_STUB_SIZE as u64,
            nt_map_view_of_section_num,
        )));

        println!(
            "[+] Syscall stubs initialized: NtOpenFile=0x{:016x}, NtCreateSection=0x{:016x}, NtMapViewOfSection=0x{:016x}",
            NT_OPEN_FILE_CLEAN.unwrap() as usize,
            NT_CREATE_SECTION_CLEAN.unwrap() as usize,
            NT_MAP_VIEW_OF_SECTION_CLEAN.unwrap() as usize
        );

        true
    }
}

fn load_ntdll_into_section() -> Result<u64> {
    unsafe {
        let mut h_file = HANDLE::default();
        let mut object_path = UNICODE_STRING::default();
        let mut object_attributes = OBJECT_ATTRIBUTES::default();
        let mut io_status_block = IO_STATUS_BLOCK::default();
        let mut h_section = HANDLE::default();
        let mut lpv_section = null_mut::<u8>();
        let mut view_size = 0usize;

        let ntdll_path = U16CString::from_str("\\??\\C:\\Windows\\System32\\ntdll.dll")
            .map_err(|_| Error::from_win32())?;

        let rtl_init_unicode_string = GetProcAddress(
            GetModuleHandleW(PCWSTR::from_raw(
                "ntdll.dll\0".encode_utf16().collect::<Vec<u16>>().as_ptr(),
            ))
            .map_err(|_| Error::from_win32())?,
            PSTR(b"RtlInitUnicodeString\0".as_ptr() as _),
        )
        .ok_or_else(|| Error::from_win32())?;

        let rtl_init_unicode_string: RtlInitUnicodeString = transmute(rtl_init_unicode_string);
        rtl_init_unicode_string(&mut object_path, PCWSTR::from_raw(ntdll_path.as_ptr()));

        object_attributes.Length = size_of::<OBJECT_ATTRIBUTES>() as u32;
        object_attributes.ObjectName = &object_path;
        object_attributes.Attributes = OBJ_CASE_INSENSITIVE;

        let status = NT_OPEN_FILE_CLEAN.unwrap()(
            &mut h_file,
            0x0001, // FILE_READ_DATA
            &object_attributes,
            &mut io_status_block,
            0x0001, // FILE_SHARE_READ
            0,
        );

        if !NT_SUCCESS(status.0) {
            #[cfg(feature = "debug")]
            println!("[!] NtOpenFile failed with status: {:#x}", status.0);
            return Err(Error::from_hresult(status.into()));
        }

        let status = NT_CREATE_SECTION_CLEAN.unwrap()(
            &mut h_section,
            0xF001F, // SECTION_ALL_ACCESS
            null(),
            null(),
            PAGE_READONLY.0,
            0x8000000, // SEC_COMMIT
            h_file,
        );

        if !NT_SUCCESS(status.0) {
            CloseHandle(h_file).map_err(|_| Error::from_win32())?;
            #[cfg(feature = "debug")]
            println!("[!] NtCreateSection failed with status: {:#x}", status.0);
            return Err(Error::from_hresult(status.into()));
        }

        let status = NT_MAP_VIEW_OF_SECTION_CLEAN.unwrap()(
            h_section,
            GetCurrentProcess(),
            &mut lpv_section,
            0,
            0,
            null(),
            &mut view_size,
            1,
            0,
            PAGE_READONLY.0,
        );

        CloseHandle(h_section).map_err(|_| Error::from_win32())?;
        CloseHandle(h_file).map_err(|_| Error::from_win32())?;

        if !NT_SUCCESS(status.0) {
            #[cfg(feature = "debug")]
            println!("[!] NtMapViewOfSection failed with status: {:#x}", status.0);
            return Err(Error::from_hresult(status.into()));
        }

        #[cfg(feature = "debug")]
        println!(
            "[+] Loaded ntdll at 0x{:x}, size 0x{:x}",
            lpv_section as u64, view_size
        );
        Ok(lpv_section as u64)
    }
}

fn rva_to_file_offset_pointer(module: u64, rva: u32) -> u64 {
    unsafe {
        let dos_header = module as *const IMAGE_DOS_HEADER;
        let nt_headers = (module + (*dos_header).e_lfanew as u64) as *const IMAGE_NT_HEADERS64;

        let section_header = &(&(*nt_headers)).OptionalHeader as *const _ as *const u8;
        let section_header = section_header
            .add((*nt_headers).FileHeader.SizeOfOptionalHeader as usize)
            as *const IMAGE_SECTION_HEADER;

        for i in 0..(*nt_headers).FileHeader.NumberOfSections {
            let section = section_header.add(i as usize);
            let section_rva = (*section).VirtualAddress as u64;
            let section_size = (*section).Misc.VirtualSize as u64;

            if rva as u64 >= section_rva && (rva as u64) < section_rva + section_size {
                let offset = rva as u64 - section_rva + (*section).PointerToRawData as u64;
                return module + offset;
            }
        }

        0
    }
}

fn find_bytes(source_addr: u64, source_length: u32) -> u64 {
    let mut source_addr = source_addr;
    let mut source_length = source_length;
    let search_bytes = [0x0f, 0x05, 0xc3]; // syscall; ret

    while source_length >= 3 {
        unsafe {
            let test = from_raw_parts(source_addr as *const u8, 3);
            if test == search_bytes {
                return source_addr;
            }
            source_addr += 1;
            source_length -= 1;
        }
    }
    0
}

fn extract_syscalls(ntdll_base: u64) -> u32 {
    unsafe {
        let dos_header = ntdll_base as *const IMAGE_DOS_HEADER;
        let nt_headers = (ntdll_base + (*dos_header).e_lfanew as u64) as *const IMAGE_NT_HEADERS64;
        let data_directory = &(*nt_headers).OptionalHeader.DataDirectory;
        let export_dir_rva = data_directory[0].VirtualAddress;
        let export_directory =
            rva_to_file_offset_pointer(ntdll_base, export_dir_rva) as *const IMAGE_EXPORT_DIRECTORY;

        if export_directory.is_null() {
            #[cfg(feature = "debug")]
            println!("[!] Failed to get export directory");
            return 0;
        }

        let number_of_names = (*export_directory).NumberOfNames;

        let functions =
            rva_to_file_offset_pointer(ntdll_base, (*export_directory).AddressOfFunctions)
                as *const u32;
        let names = rva_to_file_offset_pointer(ntdll_base, (*export_directory).AddressOfNames)
            as *const u32;
        let ordinals =
            rva_to_file_offset_pointer(ntdll_base, (*export_directory).AddressOfNameOrdinals)
                as *const u16;

        if functions.is_null() || names.is_null() || ordinals.is_null() {
            #[cfg(feature = "debug")]
            println!("[!] Invalid export table pointers");
            return 0;
        }

        let p_stubs = VirtualAlloc(
            None,
            MAX_NUMBER_OF_SYCALLS * MAX_SYSCALL_STUB_SIZE,
            MEM_RESERVE | MEM_COMMIT,
            PAGE_EXECUTE_READWRITE,
        ) as u64;

        if p_stubs == 0 {
            #[cfg(feature = "debug")]
            println!("[!] VirtualAlloc for syscall stubs failed");
            return 0;
        }

        let mut count = 0u32;
        for i in 0..number_of_names {
            if count >= MAX_NUMBER_OF_SYCALLS as u32 {
                break;
            }

            let function_name_rva = *(names.add(i as usize));
            let function_name_ptr =
                rva_to_file_offset_pointer(ntdll_base, function_name_rva) as *const i8;
            if function_name_ptr.is_null() {
                continue;
            }

            let sz_name = std::ffi::CStr::from_ptr(function_name_ptr);
            if let Ok(name) = sz_name.to_str() {
                if name.starts_with("Zw") {
                    let function_ordinal = *(ordinals.add(i as usize)) as usize;
                    let function_rva = *(functions.add(function_ordinal));
                    let function_ptr = rva_to_file_offset_pointer(ntdll_base, function_rva);
                    if function_ptr == 0 {
                        continue;
                    }

                    let function_end = find_bytes(function_ptr, MAX_SYSCALL_STUB_SIZE as u32) + 3;
                    if function_end <= 3 {
                        continue;
                    }

                    let mut syscall_name = name.to_string();
                    syscall_name.replace_range(0..2, "Nt");
                    SYSCALL_NAMES[count as usize] = syscall_name;

                    let dest_stub = p_stubs + (count as u64 * MAX_SYSCALL_STUB_SIZE as u64);
                    std::ptr::copy_nonoverlapping(
                        function_ptr as *const u8,
                        dest_stub as *mut u8,
                        (function_end - function_ptr) as usize,
                    );
                    SYSCALL_STUBS[count as usize] = dest_stub;

                    #[cfg(feature = "debug")]
                    println!(
                        "[+] Stored syscall {}: {}",
                        count, SYSCALL_NAMES[count as usize]
                    );
                    count += 1;
                }
            }
        }

        println!("[+] Extracted {} syscalls", count);
        count
    }
}

fn get_syscall(syscall_name: &str, count: u32) -> u64 {
    unsafe {
        for i in 0..count as usize {
            if SYSCALL_NAMES[i] == syscall_name {
                return SYSCALL_STUBS[i];
            }
        }
        #[cfg(feature = "debug")]
        println!("[!] Syscall not found: {}", syscall_name);
        0
    }
}

fn main() -> Result<()> {
    unsafe {
        if !init_syscalls_from_ldrp_thunk_signature() {
            #[cfg(feature = "debug")]
            println!("[!] LdrpThunk Error");
            return Err(Error::from_win32());
        }

        let p_ntdll = load_ntdll_into_section()?;
        if p_ntdll == 0 {
            #[cfg(feature = "debug")]
            println!("[!] Load Ntdll Error");
            return Err(Error::from_win32());
        }

        let ui_count = extract_syscalls(p_ntdll);
        if ui_count == 0 {
            #[cfg(feature = "debug")]
            println!("[!] Extract Syscall Error");
            return Err(Error::from_win32());
        }

        let nt_create_thread_ex_return = get_syscall("NtCreateThreadEx", ui_count);
        if nt_create_thread_ex_return == 0 {
            #[cfg(feature = "debug")]
            println!("[!] Syscall not found");
            return Err(Error::from_win32());
        }

        let nt_create_thread_ex = transmute::<u64, NtCreateThreadEx>(nt_create_thread_ex_return);
        let mut h_thread = HANDLE::default();
        let nt_status = nt_create_thread_ex(
            &mut h_thread,
            0x1FFFFF, // GENERIC_ALL
            null(),
            GetCurrentProcess(),
            0x41414141 as *mut u8, // Dummy address
            null_mut(),
            0,
            0,
            0,
            0,
            null_mut(),
        );

        if !NT_SUCCESS(nt_status.0) {
            #[cfg(feature = "debug")]
            println!(
                "[!] NtCreateThreadEx failed with status: {:#x}",
                nt_status.0
            );
            return Err(Error::from_hresult(nt_status.into()));
        }

        println!("[+] Compile Success");
        Sleep(INFINITE);

        Ok(())
    }
}
